import math
import argparse
import numpy as np

def gen_q15_rfft_h_code(n):
  h_code = '''
#ifndef __Q15_RFFT{n}_H__
#define __Q15_RFFT{n}_H__

/* below definitions are platform dependent, you may need to modify them */

typedef signed short   sint16;
typedef unsigned short uint16;
typedef signed int     sint32;
typedef unsigned int   uint32;

#define sint32_arith_shift_right(x, n) ({{ \\
  sint32 ans = (x); \\
  ans = ans >> (n); \\
  ans; }})

/* finish platform dependent definitions */


/*
 * {n} points real valued fast fourier transform
 * - x point to a buffer for both input and output.
 * - the input is real valued.
 * - the output is complex valued, only 0th~(N/2)th values are stored because
 *   the result is symmetrical, moreover, 0th and (N/2)th imag part is not stored
 *   because they are always 0, for all stored values, ith real part is stored in
 *   x[i] and ith imag part is stored in x[N-i].
 * - input and output data is Q15 fixed point format, and output data is scaled
 *   by 1.0/N.
 */
void q15_rfft{n}(sint16 *x);

#endif /* #ifndef __Q15_RFFT{n}_H__ */
'''.format(n=n)
  return h_code

def gen_q15_rfft_c_code(n):
  log2n = int(math.log2(n))
  complex_functions_definition = '''\
typedef struct {
  sint32 real;
  sint32 imag;
} complex_t;
'''

  w_table_definition = 'static const complex_t w_table_{n}[{n}/4]={{\n'.format(n = n)
  for i in range(int(n/4)):
    real = round(math.cos(2*math.pi*i/n) * 32768)
    imag = round(-math.sin(2*math.pi*i/n) * 32768)
    if real > 32767:
      real = 32767
    w_table_definition += '  {{ {real}, {imag} }}, \n'.format(real = real, imag = imag)
  w_table_definition += '};\n'

  bit_reversal_table = []
  for i in range(n):
    idx1 = i
    idx2 = int('{:0{bits}b}'.format(idx1, bits=log2n)[::-1], 2)
    if idx1 < idx2:
      bit_reversal_table.append((idx1, idx2))

  bit_reversal_table_size = len(bit_reversal_table) 
  bit_reversal_table_definition = 'static const uint16 bit_reversal_table_{n}[{bit_reversal_table_size}][2]={{\n'\
                                   .format(n = n, bit_reversal_table_size = bit_reversal_table_size)
  for idx1, idx2 in bit_reversal_table:
    bit_reversal_table_definition += '  {{ {idx1}, {idx2} }}, \n'.format(idx1 = idx1, idx2 = idx2)
  bit_reversal_table_definition += '};\n'

  c_code = '''\
#include "q15_rfft{n}.h"

{complex_functions_definition}

{w_table_definition}

{bit_reversal_table_definition}

void q15_rfft{n}(sint16 *x) {{
  uint16 stage, block, i, stride, edirts, iw;

  for(i=0; i<{bit_reversal_table_size}; i++) {{
    uint16 idx1 = bit_reversal_table_{n}[i][0];
    uint16 idx2 = bit_reversal_table_{n}[i][1];
    sint16 tmp = x[idx1];
    x[idx1] = x[idx2];
    x[idx2] = tmp;
  }}

  stride = 1;
  edirts = {n}/2;

  for( stage=0; stage<{log2n}; stage++ ) {{
    for( block=0; block<{n}; block+=(2*stride)) {{
      int iw = edirts;
      {{
        sint32 p_real = x[block];
        sint32 q_real = x[block + stride];

        sint32 pp_real = sint32_arith_shift_right(p_real + q_real, 1);
        sint32 qq_real = sint32_arith_shift_right(p_real - q_real, 1);

        x[block] = pp_real;
        x[block + stride] = qq_real;
      }}

      for( i=1; i<stride/2; i++) {{
        sint32 p_real = x[block + i];
        sint32 p_imag = x[block + stride - i];
        sint32 q_real = x[block + i + stride];
        sint32 q_imag = x[block + 2*stride - i];

        p_real = sint32_arith_shift_right(p_real, 1);
        p_imag = sint32_arith_shift_right(p_imag, 1);
        q_real = sint32_arith_shift_right(q_real, 1);
        q_imag = sint32_arith_shift_right(q_imag, 1);

        sint32 w_real = w_table_{n}[iw].real;
        sint32 w_imag = w_table_{n}[iw].imag;

        sint32 wq_real = sint32_arith_shift_right(w_real * q_real - w_imag * q_imag, 15);
        sint32 wq_imag = sint32_arith_shift_right(w_real * q_imag + w_imag * q_real, 15);
        
        sint32 pp_real = p_real + wq_real;
        sint32 pp_imag = p_imag + wq_imag;
        
        sint32 qq_real = p_real - wq_real;
        sint32 qq_imag = p_imag - wq_imag;

        x[block + i] = pp_real;
        x[block + 2*stride - i] = pp_imag;
        x[block + stride - i] = qq_real;
        x[block + i + stride] = -qq_imag;

        iw += edirts;
      }}

      if(stage > 0) {{
        sint32 p_real = x[block + stride/2];
        sint32 q_real = x[block + stride/2 + stride];

        sint32 pp_real = sint32_arith_shift_right(p_real, 1);
        sint32 pp_imag = sint32_arith_shift_right(-q_real, 1);

        x[block + stride/2] = pp_real;
        x[block + stride/2 + stride] = pp_imag;
      }}

    }}
    stride <<= 1;
    edirts >>= 1;
  }}
}}
'''.format(complex_functions_definition=complex_functions_definition,
           w_table_definition=w_table_definition,
           bit_reversal_table_definition=bit_reversal_table_definition,
           n=n,
           log2n=log2n,
           bit_reversal_table_size=bit_reversal_table_size)
  return c_code

def gen_q15_rfft_test_code(n):
  #x = np.array([0.001*(i+1) * math.cos(2*math.pi*(i+1)/n) for i in range(n)])
  x = np.random.uniform(-0.9, 0.9, n)
  ref_y = np.fft.rfft(x)

  macro_definition = '''\
typedef struct {
  float real;
  float imag;
} complex_t;

#define complex_sub(a, b) ({ \\
  complex_t ans; \\
  ans.real = a.real - b.real; \\
  ans.imag = a.imag - b.imag; \\
  ans; })

#define complex_mag(c) ( sqrt(c.real * c.real +  c.imag * c.imag) )
'''

  x_definition = 'static const float x[{n}] = {{\n'.format(n=n)
  for i in range(n):
    x_definition += '  {}, \n'.format(x[i])
  x_definition += '};\n'

  ref_y_definition = 'static const float ref_y[{n}+2] = {{\n'.format(n=n)
  for i in range(int(n/2)+1):
    ref_y_definition += '  {re}, {im},\n'.format(re = ref_y[i].real, im = ref_y[i].imag)
  ref_y_definition += '};\n'

  hyp_y_definition = 'static sint16 hyp_y[{n}] = {{\n'.format(n=n)
  for i in range(n):
    hyp_y_definition += '  {},\n'.format(round(x[i]*32768))
  hyp_y_definition += '};\n'

  test_code = '''\
#include "q15_rfft{n}.h"
#include "math.h"

{macro_definition}

{x_definition}

{ref_y_definition}

{hyp_y_definition}

float q15_rfft{n}_test_error_degree() {{
  int i;
  float ref_y_max_mag=0.0;
  float y_err_max_mag=0.0;

  q15_rfft{n}(hyp_y);
  for(i=0; i<{n}/2+1; i++) {{
    float ref_y_mag, y_err_mag;
    complex_t ref_y_i = {{ ref_y[2*i], ref_y[2*i+1] }};
    complex_t hyp_y_i = {{ hyp_y[i], 0.0}};
    if((i>0) && (i<{n}/2)) hyp_y_i.imag = hyp_y[{n} - i];
    hyp_y_i.real = hyp_y_i.real *{n}/32768;
    hyp_y_i.imag = hyp_y_i.imag *{n}/32768;

    complex_t y_diff = complex_sub(ref_y_i, hyp_y_i);
    ref_y_mag = complex_mag(ref_y_i);
    y_err_mag = complex_mag(y_diff);
    if(ref_y_mag > ref_y_max_mag) ref_y_max_mag = ref_y_mag;
    if(y_err_mag > y_err_max_mag) y_err_max_mag = y_err_mag;
  }}
  return (y_err_max_mag/ref_y_max_mag);
}}

sint16 q15_rfft{n}_test_speed(int repeat) {{
  int i;
  for(i=0; i<repeat; i++) {{
    q15_rfft{n}(hyp_y);
  }}
  return hyp_y[0];
}}

'''.format(n=n,
           macro_definition=macro_definition,
           x_definition=x_definition,
           ref_y_definition=ref_y_definition,
           hyp_y_definition=hyp_y_definition)

  return test_code

def gen_rfft(n):
  h_code = gen_q15_rfft_h_code(n)
  c_code = gen_q15_rfft_c_code(n)
  test_code = gen_q15_rfft_test_code(n)
  return (h_code, c_code, test_code)

if __name__ == "__main__":
  parser = argparse.ArgumentParser()
  parser.add_argument("n", type=int)
  parser.add_argument('path')
  args = parser.parse_args()

  n = args.n
  path = args.path

  h_code, c_code, test_code = gen_rfft(n)

  h_fname = '{path}/q15_rfft{n}.h'.format(path=path, n=n)
  c_fname = '{path}/q15_rfft{n}.c'.format(path=path,n=n)
  test_fname = '{path}/q15_rfft{n}_test.c'.format(path=path,n=n)

  open(h_fname, mode='w', encoding='utf-8').write(h_code)
  open(c_fname, mode='w', encoding='utf-8').write(c_code)
  open(test_fname, mode='w', encoding='utf-8').write(test_code)
